from mindquantum.core import Circuit, Hamiltonian, UN, H, ZZ, RX, QubitOperator
from mindquantum.framework import MQAnsatzOnlyLayer
from mindquantum.simulator import Simulator
import networkx as nx
import mindspore.nn as nn
import mindspore as ms

def build_hc(g, para):
    hc = Circuit()
    for i in g.edges:
        hc += ZZ(para).on(i)
    return hc
def build_hb(g, para):
    hc = Circuit()
    for i in g.nodes:
        hc += RX(para).on(i)
    return hc
def build_ham(g):
    hc = QubitOperator()
    for i in g.edges:
        hc += QubitOperator(f'Z{i[0]} Z{i[1]}')
    return hc
def build_ansatz(g, p):                    # g是max-cut问题的图，p是ansatz线路的层数
    circ = Circuit()                       # 创建量子线路
    for i in range(p):
        circ += build_hc(g, f'g{i}')       # 添加Uc对应的线路，参数记为g0、g1、g2...
        circ += build_hb(g, f'b{i}')       # 添加Ub对应的线路，参数记为b0、b1、b2...
    return circ

def train_params(g, p, num_of_steps = 200, learning_rate=0.05):
    
    ham = Hamiltonian(build_ham(g))              # 生成哈密顿量
    init_state_circ = UN(H, g.nodes)             # 生成均匀叠加态，即对所有量子比特作用H门
    ansatz = build_ansatz(g, p)                  # 生成ansatz线路
    circ = init_state_circ + ansatz              # 将初始化线路与ansatz线路组合成一个线路

    ms.context.set_context(mode=ms.context.PYNATIVE_MODE, device_target="CPU")

    sim = Simulator('projectq', circ.n_qubits)                     # 创建模拟器，backend使用‘projectq’，能模拟5个比特（'circ'线路中包含的比特数）
    grad_ops = sim.get_expectation_with_grad(ham, circ)            # 获取计算变分量子线路的期望值和梯度的算子
    net = MQAnsatzOnlyLayer(grad_ops)                              # 生成待训练的神经网络
    opti = nn.Adam(net.trainable_params(), learning_rate)     # 设置针对网络中所有可训练参数、学习率为0.05的Adam优化器
    train_net = nn.TrainOneStepCell(net, opti)                     # 对神经网络进行一步训练

    for i in range(num_of_steps):
        cut = (len(g.edges) - train_net()) / 2      # 将神经网络训练一步并计算得到的结果（切割边数）。注意：每当'train_net()'运行一次，神经网络就训练了一步
        if i%10 == 0:
            print("train step:", i, ", cut:", cut)  # 每训练10步，打印当前训练步数和当前得到的切割边数

    optimal_params = dict(zip(ansatz.params_name, net.weight.asnumpy())) # 获取线路参数
    return optimal_params

if __name__ == '__main__':
    g = nx.random_regular_graph(4,5)
    # pylint: disable=W0104
    p = 4
    optimal_params = train_params(g, p)
    print('Optimal Parameters: ',optimal_params)
    